#include "virtio_stor.h"

#if defined(EVENT_TRACING)
#include "virtio_stor_prefetch.tmh"
#endif
#include <ntstrsafe.h>

#define PREFETCH_REGISTRY L"\\Registry\\Machine\\SYSTEM\\ControlSet001\\Services\\viostor\\Parameters"
DECLARE_CONST_UNICODE_STRING(prefetchString, PREFETCH_REGISTRY);

VOID CopyRoutine(
    IN PVOID StartContext
    );

VOID DestroyPrefetchContext(
    IN PADAPTER_EXTENSION adaptExt
    );

static PMDL AllocatePages(ULONG bytes, BOOLEAN zero_alloc)
{
    PHYSICAL_ADDRESS low, high, skip;
    ULONG flags = MM_ALLOCATE_FULLY_REQUIRED;

    low.QuadPart = 0;
    high.QuadPart = -1;
    skip.QuadPart = 0;

    if (zero_alloc == FALSE) {
        flags |= MM_DONT_ZERO_ALLOCATION;
    }

    return MmAllocatePagesForMdlEx(low, high, skip, bytes, MmCached, flags);
}

static BOOLEAN InitPrefetchRequest(PADAPTER_EXTENSION adaptExt)
{
    PPREFETCH_INFO info = &adaptExt->prefetch_info;
    for (int i = 0; i < DEFAULT_PREFETCH_IO_COUNT; ++i) {
        info->pr[i].hdrMdl = AllocatePages(PAGE_SIZE, TRUE);
        if (!info->pr[i].hdrMdl) {
            RhelDbgPrint(TRACE_LEVEL_ERROR, " AllocatePages failed for request %d's header\n", i);
            goto err;
        }

        info->pr[i].vbr = MmMapLockedPagesSpecifyCache(info->pr[i].hdrMdl, KernelMode,
            MmCached, NULL, FALSE, NormalPagePriority | MdlMappingNoExecute);
        if (!info->pr[i].vbr) {
            RhelDbgPrint(TRACE_LEVEL_ERROR, " MmMap pages failed for request %d's header\n", i);
            goto err;
        }

        info->pr[i].dataMdl = AllocatePages(ROUND_TO_PAGES(DEFAULT_PREFETCH_SIZE), FALSE);
        if (!info->pr[i].dataMdl) {
            RhelDbgPrint(TRACE_LEVEL_ERROR, " AllocatePages failed for request %d's data\n", i);
            goto err;
        }

        info->pr[i].dataPtr = MmMapLockedPagesSpecifyCache(info->pr[i].dataMdl, KernelMode,
            MmCached, NULL, FALSE, NormalPagePriority | MdlMappingNoExecute);
        if (!info->pr[i].dataPtr) {
            RhelDbgPrint(TRACE_LEVEL_ERROR, " MmMap pages failed for request %d's data\n", i);
            goto err;
        }
    }

    return TRUE;

err:
    DestroyPrefetchContext(adaptExt);
    return FALSE;
}

static BOOLEAN CreatePrefetchThread(PADAPTER_EXTENSION adaptExt)
{
    NTSTATUS status;
    HANDLE   hThread;

    KeInitializeEvent(&adaptExt->prefetch_info.copy_event, SynchronizationEvent, FALSE);
    status = PsCreateSystemThread(&hThread, THREAD_ALL_ACCESS, NULL, NULL, NULL, CopyRoutine, adaptExt);
    if (!NT_SUCCESS(status)) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, " failed to create copy_thread, status 0x%x\n", status);
        return FALSE;
    }
    ObReferenceObjectByHandle(hThread, THREAD_ALL_ACCESS, NULL, KernelMode,
        (PVOID*)&adaptExt->prefetch_info.copy_thread, NULL);
    KeSetPriorityThread(adaptExt->prefetch_info.copy_thread, LOW_REALTIME_PRIORITY);
    ZwClose(hThread);

    return TRUE;
}

static BOOLEAN PrefetchIOEnabled(VOID)
{
    NTSTATUS ntStatus;
    UNICODE_STRING regKey;
    OBJECT_ATTRIBUTES ObjectAttributes;
    HANDLE hRegistry = NULL;
    ULONG outLen;
    UCHAR regValue[sizeof(KEY_VALUE_PARTIAL_INFORMATION) + sizeof(DWORD)];
    UINT32 prefetchStatus = 0;

    InitializeObjectAttributes(&ObjectAttributes,
        (PUNICODE_STRING)&prefetchString,
        OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE,
        NULL,
        NULL);
    ntStatus = ZwOpenKey(&hRegistry, KEY_ALL_ACCESS, &ObjectAttributes);
    if (!NT_SUCCESS(ntStatus)) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, "Open prefetching register key failed 0x%08x \n", ntStatus);
        return FALSE;
    }

    RtlInitUnicodeString(&regKey, L"IoPrefetchStatus");
    ntStatus = ZwQueryValueKey(hRegistry, &regKey,
        KeyValuePartialInformation, regValue, sizeof(regValue), &outLen);
    if (!NT_SUCCESS(ntStatus)) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, "Query prefetching register key failed 0x%08x \n", ntStatus);
        ZwClose(hRegistry);
        return FALSE;
    }
    ZwClose(hRegistry);

    if (memcpy_s(&prefetchStatus, sizeof(prefetchStatus), ((PKEY_VALUE_PARTIAL_INFORMATION)regValue)->Data,
        ((PKEY_VALUE_PARTIAL_INFORMATION)regValue)->DataLength) != 0) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, "Copy prefetching registry failed\n");
        return FALSE;
    }

    return prefetchStatus != 0;
}

BOOLEAN InitPrefetchContext(IN PADAPTER_EXTENSION adaptExt)
{
    memset(&adaptExt->prefetch_info, 0, sizeof(PREFETCH_INFO));

    if (!PrefetchIOEnabled()) {
        SetPrefetchedFeatureStatus(adaptExt, DESTROYED);
        return TRUE;
    }

    KeInitializeSpinLock(&adaptExt->prefetch_info.prefetch_lock);

    if (FALSE == InitPrefetchRequest(adaptExt)) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, " Failed to initialize prefetch request\n");
        return FALSE;
    }

    if (FALSE == CreatePrefetchThread(adaptExt)) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, " Failed to create prefetch thread\n");
        DestroyPrefetchContext(adaptExt);
        return FALSE;
    }

    PPREFETCH_INFO info = &adaptExt->prefetch_info;
    for (int i = 0; i < DEFAULT_PREFETCH_IO_COUNT; ++i) {
        info->pr[i].vbr->copy_event = &info->copy_event;
    }
    SetPrefetchedFeatureStatus(adaptExt, INITIALED);

    return TRUE;
}

VOID DestroyPrefetchContext(IN PADAPTER_EXTENSION adaptExt)
{
    PPREFETCH_INFO info = &adaptExt->prefetch_info;

    if (info->copy_thread) {
        KeSetEvent(&info->copy_event, EVENT_INCREMENT, FALSE);
        NTSTATUS status = KeWaitForSingleObject(info->copy_thread, Executive, KernelMode, FALSE, NULL);
        if (!NT_SUCCESS(status)) {
            RhelDbgPrint(TRACE_LEVEL_ERROR, " Wait copy_thread exit failed, status 0x%08x\n", status);
        }
        ObDereferenceObject(info->copy_thread);
        info->copy_thread = NULL;
    }

    for (int i = 0; i < DEFAULT_PREFETCH_IO_COUNT; ++i) {
        if (info->pr[i].vbr) {
            MmUnmapLockedPages(info->pr[i].vbr, info->pr[i].hdrMdl);
            info->pr[i].vbr = NULL;
        }
        if (info->pr[i].hdrMdl) {
            MmFreePagesFromMdl(info->pr[i].hdrMdl);
            ExFreePool(info->pr[i].hdrMdl);
            info->pr[i].hdrMdl = NULL;
        }
        if (info->pr[i].dataPtr) {
            MmUnmapLockedPages(info->pr[i].dataPtr, info->pr[i].dataMdl);
            info->pr[i].dataPtr = NULL;
        }
        if (info->pr[i].dataMdl) {
            MmFreePagesFromMdl(info->pr[i].dataMdl);
            ExFreePool(info->pr[i].dataMdl);
            info->pr[i].dataMdl = NULL;
        }
    }

    SetPrefetchedFeatureStatus(adaptExt, DESTROYED);
    info->seq_end = 0;
    info->seq_start = 0;
}

/* Fill the allocated memory to scatter list for prefetching IO */
static int FillScatterGatherFromMdl(struct scatterlist sg[], PMDL mdl, size_t length)
{
    int sgCount = 0;

    while (mdl != NULL) {
        ULONG pageOffset = MmGetMdlByteOffset(mdl);
        ULONG mdlByteCount = MmGetMdlByteCount(mdl);
        ULONG totalPages = BYTES_TO_PAGES(pageOffset + mdlByteCount);
        PPFN_NUMBER pfnArray = MmGetMdlPfnArray(mdl);
        for (ULONG i = 0; i < totalPages && length > 0; i++) {
            ULONG chunkSize = PAGE_SIZE - pageOffset;
            ULONG len = (ULONG)(min(length, chunkSize));
            length -= len;

            ULONGLONG off_end = 0;
            if (sgCount > 0) {
                off_end = sg[sgCount - 1].physAddr.QuadPart + sg[sgCount - 1].length;
            }
            if (off_end != 0 && off_end == ((ULONGLONG)pfnArray[i] << PAGE_SHIFT)) {
                sg[sgCount - 1].length += len;
                continue;
            }

            sg[sgCount].physAddr.QuadPart = ((ULONGLONG)pfnArray[i] << PAGE_SHIFT) + pageOffset;
            sg[sgCount].length = len;
            pageOffset = 0;
            sgCount++;
        }
        if (length == 0) {
            break;
        }
        mdl = mdl->Next;
    }

    return sgCount;
}

static BOOLEAN BuildPrefetchIo(PPREFETCH_REQ pr)
{
    ULONG                 sgElement;
    PADAPTER_EXTENSION    adaptExt;
    ULONGLONG             lba;
    ULONGLONG             blocks;

    adaptExt = (PADAPTER_EXTENSION)pr->DeviceExtension;
    if (adaptExt->removed) {
        return FALSE;
    }

    lba = pr->uOffset >> SECTOR_SHIFT;
    blocks = (DEFAULT_PREFETCH_SIZE + adaptExt->info.blk_size - 1) / adaptExt->info.blk_size;
    if (lba > adaptExt->lastLBA) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, " %s SRB_STATUS_BAD_SRB_BLOCK_LENGTH lba = %llu lastLBA= %llu\n",
            __FUNCTION__, lba, adaptExt->lastLBA);
        return FALSE;
    }
    if ((lba + blocks) > adaptExt->lastLBA) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, " %s lba = %llu lastLBA= %llu blocks = %llu\n",
            __FUNCTION__, lba, adaptExt->lastLBA, blocks);
        return FALSE;
    }

    /* The request has already been initialized */
    if (pr->out == 1)
        return TRUE;

    ULONG PageOffset = MmGetMdlByteOffset(pr->hdrMdl);
    ULONG MdlByteCount = MmGetMdlByteCount(pr->hdrMdl);
    PPFN_NUMBER pfn_array = MmGetMdlPfnArray(pr->hdrMdl);

    pr->sg[0].physAddr.QuadPart = ((ULONGLONG)pfn_array[0] << PAGE_SHIFT) + PageOffset + offsetof(blk_req, out_hdr);
    pr->sg[0].length = sizeof(pr->vbr->out_hdr);
    sgElement = FillScatterGatherFromMdl(&pr->sg[1], pr->dataMdl, DEFAULT_PREFETCH_SIZE) + 1;

    pr->vbr->out_hdr.type = VIRTIO_BLK_T_IN;
    pr->vbr->out_hdr.ioprio = 0;
    pr->vbr->out_hdr.sector = lba;

    pr->vbr->req = (PVOID)NULL;
    pr->out = 1;
    pr->in = sgElement;
    pr->sg[sgElement].physAddr.QuadPart = pr->sg[0].physAddr.QuadPart + sizeof(pr->vbr->out_hdr);
    pr->sg[sgElement].length = sizeof(pr->vbr->status);

    return TRUE;
}

static BOOLEAN StartPrefetchIo(PPREFETCH_REQ pr)
{
    PADAPTER_EXTENSION  adaptExt = (PADAPTER_EXTENSION)pr->DeviceExtension;
    ULONG               QueueNumber = 0;
    ULONG               MessageId = QueueNumber + 1;
    BOOLEAN             result = FALSE;
    bool                notify = FALSE;
    STOR_LOCK_HANDLE    LockHandle = { 0 };
    struct virtqueue*   vq;

    if (adaptExt->reset_in_progress) {
        return FALSE;
    }

    vq = adaptExt->vq[QueueNumber];
    RhelDbgPrint(TRACE_LEVEL_VERBOSE, " %s QueueNumber 0x%x vq = %p\n", __FUNCTION__, QueueNumber, vq);

    VioStorVQLock(pr->DeviceExtension, MessageId, &LockHandle, FALSE);
    if (virtqueue_add_buf(vq,
        &pr->sg[0],
        pr->out, pr->in,
        pr->vbr, NULL, 0) >= 0) {
        notify = virtqueue_kick_prepare(vq);
        VioStorVQUnlock(pr->DeviceExtension, MessageId, &LockHandle, FALSE);
        pr->vbr->pstatus = PREFETCHING;
        result = TRUE;
    } else {
        VioStorVQUnlock(pr->DeviceExtension, MessageId, &LockHandle, FALSE);
        RhelDbgPrint(TRACE_LEVEL_ERROR, " %s can not add packet to queue %d.\n", __FUNCTION__, QueueNumber);
    }
    if (notify) {
        virtqueue_notify(vq);
    }

    if (adaptExt->num_queues > 1) {
        if (CHECKFLAG(adaptExt->perfFlags, STOR_PERF_OPTIMIZE_FOR_COMPLETION_DURING_STARTIO)) {
            VioStorCompleteRequest(pr->DeviceExtension, MessageId, FALSE);
        }
    }

    return result;
}

static BOOLEAN CreatePrefetchIoRequest(PPREFETCH_REQ pr)
{
    BOOLEAN ret = BuildPrefetchIo(pr);
    if (ret == FALSE) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, " BuildPrefetchIo failed\n");
        return ret;
    }

    ret = StartPrefetchIo(pr);
    if (ret == FALSE) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, " StartPrefetchIo failed\n");
    }

    return ret;
}

static BOOLEAN PrefetchedIOCompleted(int status)
{
    return status == PREFETCHED_SUCC || status == PREFETCHED_FAILED;
}

/*
 * According the current offset, get the next prefetching offset.
 * For example:
 *     offset = 0x30000
 *     pr[0].uOffset = 0x40000
 *     pr[1].uOffset = 0x30000
 *     pr[2].uOffset = 0x60000
 *     pr[3].uOffset = 0x50000
 *     pr[4].uOffset = 0x800000
 * The next prefetching offset should be 0x70000.
 */
static ULONGLONG GetPrefetchOffset(PPREFETCH_INFO info, ULONGLONG offset)
{
    int i = 0;

    /* Get the next adjacent IO */
    offset += DEFAULT_PREFETCH_SIZE;

    /* Find and skip the same offset, avoid prefetching it twice */
    while (i < DEFAULT_PREFETCH_IO_COUNT) {
        for (i = 0; i < DEFAULT_PREFETCH_IO_COUNT; ++i) {
            if (offset == info->pr[i].uOffset) {
                offset += DEFAULT_PREFETCH_SIZE;
                break;
            }
        }
    }

    return offset;
}

static void PrepareNextPrefetchRequest(IN PPREFETCH_INFO info, IN PPREFETCH_REQ pr)
{
    /* Preparing for next prefetching IO request */
    pr->uOffset = GetPrefetchOffset(info, pr->uOffset);
    pr->invalid = FALSE;
    pr->vbr->out_hdr.sector = pr->uOffset >> SECTOR_SHIFT;
    pr->vbr->pstatus = PREPARING;
    pr->vbr->status = INVALID_STATUS;
}

/* Copy the hitted and prefetched data to srb */
static BOOLEAN CopyPrefetchedData(IN PVOID DeviceExtension, IN PSCSI_REQUEST_BLOCK srb, IN PUCHAR pHdrBuf)
{
    PVOID va = NULL;
    ULONG ret = StorPortGetSystemAddress(DeviceExtension, srb, &va);
    if (ret != STOR_STATUS_SUCCESS) {
        RhelDbgPrint(TRACE_LEVEL_ERROR, " Failed to get system address, ret %lu\n", ret);
        return FALSE;
    }
    RtlCopyMemory(va, pHdrBuf, DEFAULT_PREFETCH_SIZE);

    return TRUE;
}

VOID CopyRoutine(IN PVOID StartContext)
{
    UCHAR status;
    PADAPTER_EXTENSION adaptExt = (PADAPTER_EXTENSION)StartContext;
    PPREFETCH_INFO info = &adaptExt->prefetch_info;

    for (;;) {
        KeWaitForSingleObject(&info->copy_event, Executive, KernelMode, FALSE, NULL);
        if (info->pf_status == DESTROYED) {
            break;
        }

        int pr_count = 0;
        PPREFETCH_REQ prs[DEFAULT_PREFETCH_IO_COUNT];
        KIRQL oldIRQL;
        KeAcquireSpinLock(&info->prefetch_lock, &oldIRQL);
        for (int i = 0; i < DEFAULT_PREFETCH_IO_COUNT; ++i) {
            PPREFETCH_REQ pr = &info->pr[i];
            prefetch_io_status pstatus = pr->vbr->pstatus;

            if (pstatus == PREFETCHING && pr->vbr->status != INVALID_STATUS) {
                pstatus = pr->vbr->status == 0 ? PREFETCHED_SUCC : PREFETCHED_FAILED;
                pr->vbr->pstatus = pstatus;
            } else {
                continue;
            }

            if (pr->srb == NULL || !PrefetchedIOCompleted(pstatus)) {
                continue;
            }

            PSCSI_REQUEST_BLOCK prb = pr->srb;
            PSRB_EXTENSION srbExt = SRB_EXTENSION(prb);
            InterlockedDecrement((LONG volatile*)&info->big_io_count);
            srbExt->vbr.is_big_io = FALSE;
            pr->vbr->pstatus = NONE;

            if (pstatus == PREFETCHED_FAILED) {
                pr->srb = NULL;
                KeReleaseSpinLock(&info->prefetch_lock, oldIRQL);
                CompleteRequestWithStatus(pr->DeviceExtension, (PSRB_TYPE)prb, SRB_STATUS_ERROR);
                KeAcquireSpinLock(&info->prefetch_lock, &oldIRQL);
                continue;
            }
            pr->coping_data = TRUE;
            /* Copy prefetched data to srb and complete the IO request */
            KeReleaseSpinLock(&info->prefetch_lock, oldIRQL);
            BOOLEAN ret = CopyPrefetchedData(pr->DeviceExtension, prb, pr->dataPtr);
            status = ret == TRUE ? SRB_STATUS_SUCCESS : SRB_STATUS_ERROR;
            CompleteRequestWithStatus(pr->DeviceExtension, (PSRB_TYPE)prb, status);
            KeAcquireSpinLock(&info->prefetch_lock, &oldIRQL);
            pr->coping_data = FALSE;
            pr->srb = NULL;

            if (info->pf_status == RUNNING &&
                InterlockedExchangeAdd((LONG volatile*)&info->write_discard_count, 0) == 0) {
                PrepareNextPrefetchRequest(info, pr);
                prs[pr_count++] = pr;
            }
        }
        KeReleaseSpinLock(&info->prefetch_lock, oldIRQL);

        for (int i = 0; i < pr_count; ++i) {
            if (!CreatePrefetchIoRequest(prs[i])) {
                prs[i]->vbr->pstatus = NONE;
            }
        }
    }
    PsTerminateSystemThread(STATUS_SUCCESS);
}

VOID TryPrefetchData(
    IN PVOID DeviceExtension,
    IN PSRB_EXTENSION srbExt,
    IN ULONGLONG offset,
    IN ULONGLONG size
    )
{
    PADAPTER_EXTENSION adaptExt = (PADAPTER_EXTENSION)DeviceExtension;

    /* For lots of concurrent IO, no need to prefetch */
    if (adaptExt->prefetch_info.big_io_count > DEFAULT_PREFETCH_IO_COUNT) {
        return;
    }

    int pr_count = 0;
    PPREFETCH_REQ prs[DEFAULT_PREFETCH_IO_COUNT];
    KIRQL oldIRQL;
    PPREFETCH_INFO info = &adaptExt->prefetch_info;
    KeAcquireSpinLock(&info->prefetch_lock, &oldIRQL);
    if (info->pf_status == INITIALED) {
        ResetPrefetchedIOStatus(adaptExt);
        SetPrefetchedFeatureStatus(adaptExt, RUNNING);
    } else if (info->pf_status == DESTROYED) {
        KeReleaseSpinLock(&info->prefetch_lock, oldIRQL);
        return;
    }

    if (info->pf_status == RUNNING && info->seq_end == offset) {
        info->seq_end = offset + DEFAULT_PREFETCH_SIZE;
        if (InterlockedExchangeAdd((LONG volatile*)&info->write_discard_count, 0) > 0 ||
            info->seq_end - info->seq_start < DEFAULT_SEQ_IO_SIZE) {
            KeReleaseSpinLock(&info->prefetch_lock, oldIRQL);
            return;
        }

        /* Sequential read trigger prefetching IO */
        srbExt->vbr.is_prefetching = TRUE;
        ULONGLONG tmp_offset = offset;
        for (int i = 0; i < DEFAULT_PREFETCH_IO_COUNT; ++i) {
            PPREFETCH_REQ pr = &info->pr[i];
            if (pr->vbr->pstatus == PREFETCHING || pr->vbr->pstatus == PREPARING ||
                pr->srb != NULL || pr->coping_data) {
                continue;
            }
            /*
             * Using the 'unused pr' to prefetch IO which uOffset is smaller then the
             * current IO offset or the uOffset is much bigger than the current IO offset.
             */
            if (offset > pr->uOffset ||
                pr->uOffset - offset > DEFAULT_PREFETCH_IO_COUNT * DEFAULT_PREFETCH_SIZE) {
                pr->DeviceExtension = DeviceExtension;
                pr->uOffset = tmp_offset;
                PrepareNextPrefetchRequest(info, pr);
                tmp_offset = pr->uOffset;
                prs[pr_count++] = pr;
            }
        }
    } else {
        info->seq_start = offset;
        info->seq_end = offset + size;
    }
    KeReleaseSpinLock(&info->prefetch_lock, oldIRQL);

    for (int i = 0; i < pr_count; ++i) {
        if (!CreatePrefetchIoRequest(prs[i])) {
            prs[i]->vbr->pstatus = NONE;
        }
    }
}

BOOLEAN CheckHitIo(IN PVOID DeviceExtension, IN PSCSI_REQUEST_BLOCK Srb)
{
    BOOLEAN ret = FALSE;
    KIRQL oldIRQL;
    PPREFETCH_REQ cur_pr = NULL;
    PADAPTER_EXTENSION adaptExt = (PADAPTER_EXTENSION)DeviceExtension;
    PSRB_EXTENSION srbExt = SRB_EXTENSION(Srb);
    PPREFETCH_INFO info = &adaptExt->prefetch_info;

    KeAcquireSpinLock(&info->prefetch_lock, &oldIRQL);
    for (int i = 0; i < DEFAULT_PREFETCH_IO_COUNT; ++i) {
        PPREFETCH_REQ pr = &info->pr[i];
        if (pr->srb || pr->invalid == TRUE ||
            (pr->vbr->pstatus != PREFETCHING && pr->vbr->pstatus != PREFETCHED_SUCC) ||
            (pr->uOffset >> SECTOR_SHIFT) != srbExt->vbr.out_hdr.sector) {
            continue;
        }

        if (pr->vbr->pstatus == PREFETCHING) {
            pr->srb = Srb;
            KeReleaseSpinLock(&info->prefetch_lock, oldIRQL);
            return TRUE;
        }

        InterlockedDecrement((LONG volatile*)&info->big_io_count);
        srbExt->vbr.is_big_io = FALSE;

        pr->coping_data = TRUE;
        /* The IO has been prefetched success */
        KeReleaseSpinLock(&info->prefetch_lock, oldIRQL);
        ret = CopyPrefetchedData(DeviceExtension, Srb, pr->dataPtr);
        KeAcquireSpinLock(&info->prefetch_lock, &oldIRQL);
        pr->coping_data = FALSE;
        if (ret) {
            if (InterlockedExchangeAdd((LONG volatile*)&info->write_discard_count, 0) == 0) {
                PrepareNextPrefetchRequest(info, pr);
                cur_pr = pr;
            }
        } else {
            RhelDbgPrint(TRACE_LEVEL_ERROR, " %s: copy prefetched data failed\n", __FUNCTION__);
        }
        break;
    }
    KeReleaseSpinLock(&info->prefetch_lock, oldIRQL);

    if (ret) {
        CompleteRequestWithStatus(DeviceExtension, (PSRB_TYPE)Srb, SRB_STATUS_SUCCESS);
    }

    if (cur_pr && !CreatePrefetchIoRequest(cur_pr)) {
        cur_pr->vbr->pstatus = NONE;
    }

    return ret;
}

VOID InvalidPrefetchedIO(IN PADAPTER_EXTENSION adaptExt, IN ULONGLONG offset, IN ULONGLONG length, IN BOOLEAN add_count)
{
    if (!IsPrefetchedFeatureEnabled(adaptExt)) {
        return;
    }

    KIRQL oldIRQL;
    KeAcquireSpinLock(&adaptExt->prefetch_info.prefetch_lock, &oldIRQL);
    if (add_count) {
        InterlockedIncrement((LONG volatile*)&adaptExt->prefetch_info.write_discard_count);
    }
    for (int i = 0; i < DEFAULT_PREFETCH_IO_COUNT; ++i) {
        PPREFETCH_REQ pr = &adaptExt->prefetch_info.pr[i];
        if (!pr->invalid && offset < pr->uOffset + DEFAULT_PREFETCH_SIZE && offset + length > pr->uOffset) {
            if (pr->srb == NULL) {
                pr->invalid = TRUE;
            } else {
                RhelDbgPrint(TRACE_LEVEL_WARNING, " %s: [0x%llx, 0x%llx] overlap with pr offset 0x%llx, srb is 0x%p\n",
                    __FUNCTION__, offset, offset + length, pr->uOffset, pr->srb);
            }
        }
    }
    KeReleaseSpinLock(&adaptExt->prefetch_info.prefetch_lock, oldIRQL);
}

VOID ResetPrefetchedIOStatus(IN PADAPTER_EXTENSION adaptExt)
{
    for (int i = 0; i < DEFAULT_PREFETCH_IO_COUNT; ++i) {
        adaptExt->prefetch_info.pr[i].vbr->pstatus = NONE;
    }
}

BOOLEAN IsPrefetchedFeatureEnabled(IN PADAPTER_EXTENSION adaptExt)
{
    return InterlockedExchangeAdd(&adaptExt->prefetch_info.pf_status, 0) != DESTROYED;
}

VOID SetPrefetchedFeatureStatus(IN PADAPTER_EXTENSION adaptExt, IN prefetch_feature_status status)
{
    /*
     * DESTROYED means alloc memory failed, prefetching feature is not enabled.
     * So, the status should not be changed.
     */
    if (!IsPrefetchedFeatureEnabled(adaptExt)) {
        return;
    }
    InterlockedExchange(&adaptExt->prefetch_info.pf_status, status);
}

VOID CompletePrefetchRequest(IN PADAPTER_EXTENSION adaptExt, IN PVOID DeviceExtension)
{
    KIRQL oldIRQL;
    PSCSI_REQUEST_BLOCK srbs[DEFAULT_PREFETCH_IO_COUNT];
    int i, count = 0;
    PPREFETCH_INFO info = &adaptExt->prefetch_info;

    if (!IsPrefetchedFeatureEnabled(adaptExt)) {
        return;
    }

    KeAcquireSpinLock(&info->prefetch_lock, &oldIRQL);
    for (i = 0; i < DEFAULT_PREFETCH_IO_COUNT; ++i) {
        if (NULL != info->pr[i].srb) {
            srbs[count++] = info->pr[i].srb;
            info->pr[i].srb = NULL;
        }
    }
    KeReleaseSpinLock(&info->prefetch_lock, oldIRQL);

    for (i = 0; i < count; ++i) {
        CompleteRequestWithStatus(DeviceExtension, (PSRB_TYPE)srbs[i], SRB_STATUS_BUS_RESET);
    }
}

VOID HandleCompletedPrefetchIO(IN PADAPTER_EXTENSION adaptExt, IN pblk_req vbr)
{
    if (!IsPrefetchedFeatureEnabled(adaptExt)) {
        return;
    }

    if (vbr->is_big_io) {
        InterlockedDecrement((LONG volatile*)&adaptExt->prefetch_info.big_io_count);
        vbr->is_big_io = FALSE;
    }
    if (NULL == vbr->req && vbr->pstatus == PREFETCHING) {
        KeSetEvent(vbr->copy_event, EVENT_INCREMENT, FALSE);
        return;
    }
    if (vbr->out_hdr.type == VIRTIO_BLK_T_OUT || vbr->out_hdr.type == VIRTIO_BLK_T_DISCARD) {
        InterlockedDecrement((LONG volatile*)&adaptExt->prefetch_info.write_discard_count);
    }
}
